%% On the Optimal Design of Transfers and Income-Tax Progressivity
% Model equivalent of Golosov, Graber, Mogstad, Novgorodsky (2021)

clear; clc; close all;

%% Read model outputs

% State space
na = 300;
nx = 19;
nS = 300*19;

% Other parameters
alpha = 0.64;
delta = 0.06;

% Asset grid
a_vec = dlmread('a_vec.txt')';
a_vec = a_vec(:);

% Productivity grid
x_vec = dlmread('x_vec.txt')';
x_vec = x_vec(:);
x_vec = x_vec(1:nx);

% Labor policy
h_pol = dlmread('h_pol_init.txt');
h_pol = reshape(h_pol',numel(h_pol),1);   
h_pol = reshape(h_pol(1:nS),nS,1);
h_pol = reshape(h_pol,na,nx);

% Asset policy
a_pol = dlmread('a_pol_init.txt');
a_pol = reshape(a_pol',numel(h_pol),1);   
a_pol = reshape(a_pol(1:nS),nS,1);
a_pol = reshape(a_pol,na,nx);

% Measure
mu = dlmread('mu_init.txt');
mu = reshape(mu',numel(mu),1);   
mu = reshape(mu(1:nS),nS,1);
mu = reshape(mu,na,nx);

% Interest rate
prices_INIT = dlmread('Prices_INIT.txt');
r_init = prices_INIT(1,2);
wge_init = (1-alpha)./(r_init+delta);
wge_init = alpha*(wge_init.^((1-alpha)/alpha));

% Transition matrix
pi = dlmread('Px.txt');
pi = pi';
pi = pi(1:nx*nx);
pi = reshape(pi,nx,nx);

% Labor income distribution
x_mat = repmat(x_vec',na,1);
yl = h_pol.*x_mat*wge_init;

% Total income distribution
a_mat = repmat(a_vec,1,nx);
yk = a_mat*r_init;

ytot = yk+yl;
ymean = sum(sum(mu.*ytot));


%% Forward simulation baseline

% first two indices: assets and productivities in in the pre shock distribution
% 6 time periods after that: wealth increase and 5 years after that
% in each of the 6 periods, there's a distribution over assets and productivities for each initial state
mu_simul = zeros(na,nx,6,na,nx);

% pre period is year -1
% this is year 0: in the baseline there is no wealth change; new wealth is just given by policy function
for ia = 1:na       % for each asset level in the pre-win distribution
    for ix = 1:nx   % for each productivity in the pre-win distribution
        anew = a_pol(ia,ix);                                            % new asset level in period 0 according to policy function
        aloc = find(a_vec>=anew,1);                                     % asset index above new asset level
        da = (anew-a_vec(aloc-1))/(a_vec(aloc)-a_vec(aloc-1));          % variable to distribute mass to grid points below and above asset choice
        for ixp = 1:nx                                                  % loop over all possible period 0 productivity states, given pre-win productivity
            mu_simul(ia,ix,1,aloc-1,ixp) = mu(ia,ix)*(1-da)*pi(ix,ixp); % allocate mass to point below in year 0 (indexed by 1 in the mu_simul matrix)
            mu_simul(ia,ix,1,aloc,ixp) = mu(ia,ix)*da*pi(ix,ixp);       % allocate mass to point above in year 0 (indexed by 1 in the mu_simul matrix)
        end
    end
end

% now deal with years 1 to 5
% separately for all year -1 (a,x) combinations (first two loops)
% in any of periods 0 to 5 they can be in any (a,x) state (last two loops)
for ia = 1:na                   % initial asset index (period -1)
    for ix = 1:nx               % initial productivity index (period -1)
        for t = 2:6             % 5 years: years 1 to 5 (indexed by 2 to 6, as 0 is also in this matrix, see above)
            for iap = 1:na      % for each asset level in the last period
                for ixp = 1:nx  % for each productivity level in the last period
                    anew = a_pol(iap,ixp);                                  % new asset level according to policy function
                    aloc = find(a_vec>=anew,1);                             % asset index above new asset level
                    da = (anew-a_vec(aloc-1))/(a_vec(aloc)-a_vec(aloc-1));  % variable to distribute mass to grid points below and above asset choice
                    for ixpp = 1:nx                                         % loop over all possible period 0 productivity states, given pre-win productivity
                        mu_simul(ia,ix,t,aloc-1,ixpp) = mu_simul(ia,ix,t,aloc-1,ixpp) + mu_simul(ia,ix,t-1,iap,ixp)*(1-da)*pi(ixp,ixpp);    % allocate mass to point below
                        mu_simul(ia,ix,t,aloc,ixpp) = mu_simul(ia,ix,t,aloc,ixpp) + mu_simul(ia,ix,t-1,iap,ixp)*da*pi(ixp,ixpp);            % allocate mass to point above
                    end
                end
            end
        end
    end
end

% Average income in years 1-6 for each initial state combination
yl_avg = zeros(na,nx,6);
for ia = 1:na           % initial asset level
    for ix = 1:nx       % intial productivity level
        for t = 1:6
            yl_avg(ia,ix,t) = sum(sum(squeeze(mu_simul(ia,ix,t,:,:)).*yl))/sum(sum(squeeze(mu_simul(ia,ix,t,:,:)))); % average income of people with initial state (ia,ix) in year 0 to 5 (indexed by 1:6)
        end
    end
end
% Reshape to vector
yl_avg_vec = zeros(nS,6);
for t = 1:6
    yl_aux = squeeze(yl_avg(:,:,t));
    yl_avg_vec(:,t) = yl_aux(:);
end


%% Forward simulation wealth shock

dollar_change = 180000/1.0314;    % wealth shock [180,000 is average win size; small: 50,000; medium: 500,000; large: 2,000,000]; inflation adjusted
ymean_dollar = 90185;             % mean total income in the data

% simulation matrix for case where we change wealth
mu_simul_shock = zeros(na,nx,6,na,nx);

% now new wealth level in period 0 is not just given by policy function but also by added wealth
for ia = 1:na
    for ix = 1:nx
        anew = min(a_pol(ia,ix)+dollar_change/ymean_dollar*ymean,max(a_vec));
        aloc = find(a_vec>=anew,1);
        da = (anew-a_vec(aloc-1))/(a_vec(aloc)-a_vec(aloc-1));
        for ixp = 1:nx
            mu_simul_shock(ia,ix,1,aloc-1,ixp) = mu(ia,ix)*(1-da)*pi(ix,ixp);
            mu_simul_shock(ia,ix,1,aloc,ixp) = mu(ia,ix)*da*pi(ix,ixp);
        end
    end
end

% no new shock afterwards, so this is the same as above for years 1-5 (indexed by 2 to 6)
for ia = 1:na
    for ix = 1:nx
        for t = 2:6
            for iap = 1:na
                for ixp = 1:nx
                    anew = a_pol(iap,ixp);
                    aloc = find(a_vec>=anew,1);
                    da = (anew-a_vec(aloc-1))/(a_vec(aloc)-a_vec(aloc-1));
                    for ixpp = 1:nx
                        mu_simul_shock(ia,ix,t,aloc-1,ixpp) = mu_simul_shock(ia,ix,t,aloc-1,ixpp) + mu_simul_shock(ia,ix,t-1,iap,ixp)*(1-da)*pi(ixp,ixpp);
                        mu_simul_shock(ia,ix,t,aloc,ixpp) = mu_simul_shock(ia,ix,t,aloc,ixpp) + mu_simul_shock(ia,ix,t-1,iap,ixp)*da*pi(ixp,ixpp);
                    end
                end
            end
        end
    end
end

% Average income in years 1-6 for each initial state combination
yl_avg_shock = zeros(na,nx,6);
for ia = 1:na
    for ix = 1:nx
        for t = 1:6
            yl_avg_shock(ia,ix,t) = sum(sum(squeeze(mu_simul_shock(ia,ix,t,:,:)).*yl))/sum(sum(squeeze(mu_simul_shock(ia,ix,t,:,:))));
        end
    end
end
% Reshape to vector
yl_avg_vec_shock = zeros(nS,6);
for t = 1:6
    yl_aux = squeeze(yl_avg_shock(:,:,t));
    yl_avg_vec_shock(:,t) = yl_aux(:);
end

%% Compute wealth effects

% Labor income change (comparing years 1-5 for case with and without added wealth)
yl_change = yl_avg_vec_shock(:,2:6)-yl_avg_vec(:,2:6);
yl_change_dollar = yl_change/ymean*ymean_dollar;

% Average change
yl_vec = yl(:);                             % reshape initial income distribution as changes above
mu_vec = mu(:);                             % reshape initial distribution also to vector
yl_change_dollar_mean = zeros(1,5);         % average dollar change for each year -1 (a,x) combination in years 1-5
yl_change_dollar_mean_norm = zeros(1,5);    % average dollar change for each year -1 (a,x) combination in years 1-5: normalized in $100 terms
for t = 1:5
    yl_change_dollar_mean(t) = sum(yl_change_dollar(:,t).*mu_vec);                  % does not have to be divided by sum(mu_vec) as this equals 1
    yl_change_dollar_mean_norm(t) = yl_change_dollar_mean(t)/dollar_change*100;     % normalize change
end

% Average change across income distribution, where quartiles are computed based on the t = -1 income distribution
[yl_sorted,ind] = sort(yl_vec);                                 % sort initial income distribution
mu_sorted = mu(ind);                                            % sort measure accordingly
yl_change_dollar_sorted = zeros(nS,5);                          % sort columns of change matrix accordingly
for t = 1:5
    yl_change_dollar_sorted(:,t) = yl_change_dollar(ind,t);
end
mu_cum = cumsum(mu_sorted);                     % cumulative measure
loc25 = find(mu_cum>=0.25,1);                   % index of 25th percentile
loc50 = find(mu_cum>=0.50,1);                   % index of 50th percentile
loc75 = find(mu_cum>=0.75,1);                   % index of 75th percentile
yl_change_dollar_mean_q1 = zeros(1,5);          % average income change Q1
yl_change_dollar_mean_q2 = zeros(1,5);          % average income change Q2
yl_change_dollar_mean_q3 = zeros(1,5);          % average income change Q3
yl_change_dollar_mean_q4 = zeros(1,5);          % average income change Q4
yl_change_dollar_mean_norm_q1 = zeros(1,5);     % average income change Q1: normalized in 100 dollar terms
yl_change_dollar_mean_norm_q2 = zeros(1,5);     % average income change Q2: normalized in 100 dollar terms
yl_change_dollar_mean_norm_q3 = zeros(1,5);     % average income change Q3: normalized in 100 dollar terms
yl_change_dollar_mean_norm_q4 = zeros(1,5);     % average income change Q4: normalized in 100 dollar terms
for i = 1:5
    yl_change_dollar_mean_q1(i) = sum(yl_change_dollar_sorted(1:loc25,i).*mu_sorted(1:loc25))/sum(mu_sorted(1:loc25));
    yl_change_dollar_mean_q2(i) = sum(yl_change_dollar_sorted(loc25+1:loc50,i).*mu_sorted(loc25+1:loc50))/sum(mu_sorted(loc25+1:loc50));
    yl_change_dollar_mean_q3(i) = sum(yl_change_dollar_sorted(loc50+1:loc75,i).*mu_sorted(loc50+1:loc75))/sum(mu_sorted(loc50+1:loc75));
    yl_change_dollar_mean_q4(i) = sum(yl_change_dollar_sorted(loc75+1:end,i).*mu_sorted(loc75+1:end))/sum(mu_sorted(loc75+1:end));
    yl_change_dollar_mean_norm_q1(i) = yl_change_dollar_mean_q1(i)/dollar_change*100;
    yl_change_dollar_mean_norm_q2(i) = yl_change_dollar_mean_q2(i)/dollar_change*100;
    yl_change_dollar_mean_norm_q3(i) = yl_change_dollar_mean_q3(i)/dollar_change*100;
    yl_change_dollar_mean_norm_q4(i) = yl_change_dollar_mean_q4(i)/dollar_change*100;
end

disp('Five year average')
disp(['Full sample: ', num2str(mean(yl_change_dollar_mean_norm))])
disp(['Q1: ', num2str(mean(yl_change_dollar_mean_norm_q1))])
disp(['Q2: ', num2str(mean(yl_change_dollar_mean_norm_q2))])
disp(['Q3: ', num2str(mean(yl_change_dollar_mean_norm_q3))])
disp(['Q4: ', num2str(mean(yl_change_dollar_mean_norm_q4))])
disp('Short run')
disp(['Full sample: ', num2str(mean(yl_change_dollar_mean_norm(1:2)))])
disp(['Q1: ', num2str(mean(yl_change_dollar_mean_norm_q1(1:2)))])
disp(['Q2: ', num2str(mean(yl_change_dollar_mean_norm_q2(1:2)))])
disp(['Q3: ', num2str(mean(yl_change_dollar_mean_norm_q3(1:2)))])
disp(['Q4: ', num2str(mean(yl_change_dollar_mean_norm_q4(1:2)))])
disp('Long run')
disp(['Full sample: ', num2str(mean(yl_change_dollar_mean_norm(3:5)))])
disp(['Q1: ', num2str(mean(yl_change_dollar_mean_norm_q1(3:5)))])
disp(['Q2: ', num2str(mean(yl_change_dollar_mean_norm_q2(3:5)))])
disp(['Q3: ', num2str(mean(yl_change_dollar_mean_norm_q3(3:5)))])
disp(['Q4: ', num2str(mean(yl_change_dollar_mean_norm_q4(3:5)))])

